from data_loader.simple_mnist_data_loader import SimpleMnistDataLoader
from models.simple_mnist_model import SimpleMnistModel
from trainers.simple_mnist_trainer import SimpleMnistModelTrainer
from utils.config import process_config
from utils.dirs import create_dirs
from utils.utils import get_args

def main():
    print('[INFO] 解析配置...')
    try:
        args, parser = get_args()
        config = process_config(args.config)
    except Exception as e:
        print('[Exception] 配置无效, %s' % e)
        if parser:
            parser.print_help()
        print('[Exception] 参考: python main_train.py -c configs/simple_mnist_config.json')
        exit(0)

    # create the experiments dirs
    create_dirs([config.callbacks.tensorboard_log_dir, config.callbacks.checkpoint_dir])

    print('[INFO] 加载数据...')
    data_loader = SimpleMnistDataLoader(config)

    print('[INFO] 构造网络...')
    model = SimpleMnistModel(config)

    print('[INFO] 加载训练...')
    trainer = SimpleMnistModelTrainer(model.model, data_loader.get_train_data(), config)

    print('[INFO] 开始训练...')
    trainer.train()

    print('[INFO] 训练结束...')


if __name__ == '__main__':
    main()
